{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cSrJYrFrY2aj"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "i1PHqD-ZY4-c"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YNDq8NbCY7oh"
      },
      "source": [
        "# Workshop: How to Fine-tuning Gemma\n",
        "\n",
        "To illustrate fine-tuning the model for a specific task, You'll learn how to condition a Gemma model to answer in a specific language. Let's consider the example of generating a random Portuguese title based on a user's instruction such as \"Write a title\". To make this possible, you will curate a manageable dataset that can be manually processed. This approach is feasible because Gemma 2 has prior knowledge of general Portuguese language patterns, enabling it to adapt to this specific task effectively.\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Workshops/Workshop_How_to_Fine_tuning_Gemma.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4EM3g9u2_KA"
      },
      "source": [
        "## What is Fine-tuning\n",
        "\n",
        "In the first place, you have to understand what is fine-tuning. It's a specialized form of [transfer learning](https://en.wikipedia.org/wiki/Transfer_learning). It involves taking a pre-trained language model - one that has already been exposed to a vast corpus of text data and learned the general patterns and structures of language - and further training it on a smaller, more specific dataset. This additional training allows the model to adapt and refine its knowledge, making it better suited for a particular task or domain.\n",
        "\n",
        "Imagine you are a skilled gamer who excels at various genres, from action-adventures to strategy games. Fine-tuning is akin to taking you and having you focus intensely on mastering a specific game, like a complex real-time strategy (RTS) title. You already possess a strong foundation of gaming skills and knowledge, but the dedicated practice and study within the RTS genre sharpens your tactics, understanding of game mechanics, and overall proficiency within that particular realm.\n",
        "\n",
        "Similarly, pre-trained language models have a broad understanding of language, but fine-tuning helps them specialize. By exposing them to a curated dataset relevant to your desired application, you guide the model to learn the nuances and intricacies specific to that domain. It's like giving the model a crash course in the language of your chosen field, enabling it to perform tasks with greater accuracy and fluency.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3rzH5Ugf5RlJ"
      },
      "source": [
        "## Setup\n",
        "\n",
        "### Select the Colab runtime\n",
        "To complete this tutorial, you'll need to have a Colab runtime with sufficient resources to run the Gemma model:\n",
        "\n",
        "1. In the upper-right of the Colab window, select **▾ (Additional connection options)**.\n",
        "2. Select **Change runtime type**.\n",
        "3. Under **Hardware accelerator**, select **T4 GPU**.\n",
        "\n",
        "\n",
        "### Gemma setup on Kaggle\n",
        "To complete this tutorial, you'll first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
        "\n",
        "* Get access to Gemma on kaggle.com.\n",
        "* Select a Colab runtime with sufficient resources to run the Gemma 2B model.\n",
        "* Generate and configure a Kaggle username and API key.\n",
        "\n",
        "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "URMuBzkMVxpU"
      },
      "source": [
        "### Set environment variables\n",
        "\n",
        "Set environment variables for ```KAGGLE_USERNAME``` and ```KAGGLE_KEY```."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IUOX2hqjV7Ku"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata, drive\n",
        "\n",
        "# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env\n",
        "# vars as appropriate for your system.\n",
        "os.environ[\"KAGGLE_USERNAME\"] = userdata.get(\"KAGGLE_USERNAME\")\n",
        "os.environ[\"KAGGLE_KEY\"] = userdata.get(\"KAGGLE_KEY\")\n",
        "\n",
        "# Mounting gDrive for to store artifacts\n",
        "drive.mount(\"/content/drive\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXfDwRTQVns2"
      },
      "source": [
        "### Install dependencies\n",
        "\n",
        "Install Keras and KerasNLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zHs7wpZusEML"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U keras-nlp datasets\n",
        "!pip install -q -U keras\n",
        "\n",
        "# Set the backbend before importing Keras\n",
        "os.environ[\"KERAS_BACKEND\"] = \"jax\"\n",
        "# Avoid memory fragmentation on JAX backend.\n",
        "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"1.00\"\n",
        "\n",
        "import keras_nlp\n",
        "import keras\n",
        "\n",
        "# Run at half precision.\n",
        "#keras.config.set_floatx(\"bfloat16\")\n",
        "\n",
        "# Training Configurations\n",
        "token_limit = 128\n",
        "num_data_limit = 100\n",
        "lora_name = \"my_lora\"\n",
        "lora_rank = 4\n",
        "lr_value = 1e-3\n",
        "train_epoch = 5\n",
        "model_id = \"gemma2_instruct_2b_en\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kUl0t469YfQY"
      },
      "source": [
        "## Load Model\n",
        "\n",
        "**Why Fine-tuning?**\n",
        "\n",
        "Before embarking on fine-tuning, it's crucial to evaluate if its benefits align with the specific requirements of your application. Fine-tuning involves meticulous data preparation and extensive training, making it an arduous process. Therefore, it's essential to assess whether the potential gains justify the significant effort required.\n",
        "\n",
        "**Try \"Prompt Engineering\" first.** before fine-tuning\n",
        "\n",
        "Would you like to enable Gemma's multilingual capabilities?\n",
        "Please note that Gemma 2 already has some multilingual capabilities. Here's the example output from Gemma 2 2B instruction-tuned model.\n",
        "\n",
        "Do you wish to adjust the tone or writing style?\n",
        "Gemma 2 might be familiar with the writing style you have in mind. Here's another output from the same model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gm4jIEqmYfQY"
      },
      "outputs": [],
      "source": [
        "import keras\n",
        "import keras_nlp\n",
        "\n",
        "import time\n",
        "\n",
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(model_id)\n",
        "gemma_lm.summary()\n",
        "\n",
        "tick_start = 0\n",
        "\n",
        "def tick():\n",
        "    global tick_start\n",
        "    tick_start = time.time()\n",
        "\n",
        "def tock():\n",
        "    print(f\"TOTAL TIME ELAPSED: {time.time() - tick_start:.2f}s\")\n",
        "\n",
        "def text_gen(prompt):\n",
        "    tick()\n",
        "    input = f\"<start_of_turn>user\\n{prompt}<end_of_turn>\\n<start_of_turn>model\\n\"\n",
        "    output = gemma_lm.generate(input, max_length=token_limit)\n",
        "    print(\"\\nGemma output:\")\n",
        "    print(output)\n",
        "    tock()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o3a-0aezxZ1D"
      },
      "outputs": [],
      "source": [
        "# inference before fine-tuning\n",
        "text_gen(\"Translate the text below to Portuguese.\\n\\\"Hi, how can I get to the MASP museum?\\\"\")\n",
        "text_gen(\"Speak like a pirate. Teach me why the earth is flat.\")\n",
        "text_gen(\"Write a title\")\n",
        "text_gen(\"Write a poem\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kpubX_7zUSDc"
      },
      "source": [
        "## What is a Tokenizer\n",
        "\n",
        "A tokenizer is a crucial component that bridges the gap between human language and the numerical representations that AI models can understand. It's essentially a tool that breaks down text into smaller units, called tokens, which can then be processed by the model.\n",
        "\n",
        "Think of it like this:\n",
        "\n",
        "- Humans: Communicate using words, sentences, and paragraphs.\n",
        "- LLMs: They \"think\" in terms of numbers and vectors.\n",
        "\n",
        "The tokenizer acts as a translator between these two worlds. It takes human-readable text and converts it into a format that the LLM can understand and work with.\n",
        "\n",
        "\n",
        "\n",
        "Gemma tokenizer is based on [SentencePiece](https://github.com/google/sentencepiece). The size of the vocabulary is predetermined before training. SentencePiece then learns the optimal subword segmentation based on the chosen vocabulary size and the training data.\n",
        "\n",
        "Gemma's large 256k vocabulary allows it to handle diverse text inputs and potentially improve performance on various tasks, e.g. handling multilingual text inputs.\n",
        "\n",
        "Gemma's tokenizer uses byte-level encoding, which means it can handle any character from any language, even those with different writing systems (like Chinese, Japanese, or Korean).\n",
        "\n",
        "This avoids the limitations of character-based tokenizers, which might struggle with languages that have a vast number of characters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xR1riO7NURFF"
      },
      "outputs": [],
      "source": [
        "tokenizer = keras_nlp.models.GemmaTokenizer.from_preset(model_id)\n",
        "import jax\n",
        "\n",
        "def detoken(tokens):\n",
        "  print(tokens)\n",
        "  for x in tokens:\n",
        "    word = tokenizer.detokenize(jax.numpy.array([x]))\n",
        "    print(f\"{x:6} -> {word}\")\n",
        "\n",
        "# (example text: “Hi, Nice to meet you. The weather is really nice today.”)\n",
        "detoken(tokenizer(\"olá. Prazer em conhecê-lo. O tempo está muito bom hoje.\"))\n",
        "print()\n",
        "detoken(tokenizer(\"नमस्ते। आपसे मिलकर अच्छा लगा। आज मौसम सचमुच अच्छा है.\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9T7xe_jzslv4"
      },
      "source": [
        "## Load Dataset\n",
        "\n",
        "How many datasets do you need? You can start with a relatively small dataset, approximately 10 to 20, those can have a significant impact on a model's behavior.\n",
        "\n",
        "To improve the output quality, a target of around 200 total examples is recommended. Nevertheless, the amount of data required for tuning really depends on how much you want to influence the model's behavior. Our recommendation is to commence with a limited amount of data and gradually incorporate additional data into the training process until the desired behavior is achieved."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZiS-KU9osh_N"
      },
      "outputs": [],
      "source": [
        "tokenizer = keras_nlp.models.GemmaTokenizer.from_preset(model_id)\n",
        "\n",
        "# example titles\n",
        "data = [\n",
        "    \"O Alquimista\",  # by Paulo Coelho\n",
        "    \"Dom Casmurro\",  # by Machado de Assis\n",
        "    \"Memorial do Convento\",  # by José Saramago\n",
        "    \"A Hora da Estrela\",  # by Clarice Lispector\n",
        "    \"Vidas Secas\",  # by Graciliano Ramos\n",
        "    \"O Cortiço\",  # by Aluísio Azevedo\n",
        "    \"Grande Sertão: Veredas\",  # by Guimarães Rosa\n",
        "    \"Capitães da Areia\",  # by Jorge Amado\n",
        "    \"A Sibila\",  # by Agustina Bessa-Luís\n",
        "    \"Os Maias\",  # by Eça de Queirós\n",
        "    \"O Crime do Padre Amaro\",  # by Eça de Queirós\n",
        "    \"A Relíquia\",  # by Eça de Queirós\n",
        "    \"O Primo Basílio\",  # by Eça de Queirós\n",
        "    \"A Ilustre Casa de Ramires\",  # by Eça de Queirós\n",
        "    \"A Cidade e as Serras\"  # by Eça de Queirós\n",
        "]\n",
        "\n",
        "train = []\n",
        "\n",
        "for x in data:\n",
        "  item = f\"<start_of_turn>user\\nWrite a title<end_of_turn>\\n<start_of_turn>model\\n{x}<end_of_turn>\"\n",
        "  length = len(tokenizer(item))\n",
        "  # skip data if the token length is longer than our limit\n",
        "  if length < token_limit:\n",
        "    train.append(item)\n",
        "    if(len(train)>=num_data_limit):\n",
        "      break\n",
        "\n",
        "print(len(train))\n",
        "print(train[0])\n",
        "print()\n",
        "print(train[1])\n",
        "print()\n",
        "print(train[2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9s1o96HRtwV_"
      },
      "source": [
        "See below example code, using HF datasets, if your datasets are much bigger."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BxSqXLeCt9Uu"
      },
      "outputs": [],
      "source": [
        "# Example of using HF datasets\n",
        "# Replying to request emails that a bakery business might get, in Korean\n",
        "# The given prompt means \"Please write an email reply with below.\"\n",
        "'''\n",
        "tokenizer = keras_nlp.models.GemmaTokenizer.from_preset(model_id)\n",
        "\n",
        "# prompt structure\n",
        "# <start_of_turn>user\n",
        "# 다음에 대한 이메일 답장을 작성해줘.\n",
        "# \"{EMAIL CONTENT FROM THE CUSTOMER}\"\n",
        "# <end_of_turn>\n",
        "# <start_of_turn>model\n",
        "# {MODEL ANSWER}<end_of_turn>\n",
        "\n",
        "# input, output\n",
        "from datasets import load_dataset\n",
        "ds = load_dataset(\n",
        "    \"bebechien/korean_cake_boss\",\n",
        "    split=\"train\",\n",
        ")\n",
        "print(ds)\n",
        "data = ds.with_format(\"np\", columns=[\"input\", \"output\"], output_all_columns=False)\n",
        "train = []\n",
        "\n",
        "for x in data:\n",
        "  item = f\"<start_of_turn>user\\n다음에 대한 이메일 답장을 작성해줘.\\n\\\"{x['input']}\\\"<end_of_turn>\\n<start_of_turn>model\\n{x['output']}<end_of_turn>\"\n",
        "  length = len(tokenizer(item))\n",
        "  # skip data if the token length is longer than our limit\n",
        "  if length < token_limit:\n",
        "    train.append(item)\n",
        "    if(len(train)>=num_data_limit):\n",
        "      break\n",
        "\n",
        "print(len(train))\n",
        "print(train[0])\n",
        "print(train[1])\n",
        "print(train[2])\n",
        "'''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5NTIrFbJ3dBv"
      },
      "source": [
        "In the context of a small dataset, the primary concern is that the model may prioritize memorizing specific examples rather than generalizing well to new and unobserved data. This limitation highlights the importance of utilizing a larger dataset during fine-tuning, as it enhances the model's ability to capture broader patterns and relationships."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "th0WS33gayn9"
      },
      "source": [
        "## LoRA Fine-tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ugc2ub4nau1j"
      },
      "source": [
        "![lora.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pt7Nr6a7tItO"
      },
      "source": [
        "Fine-tuning a model involves updating its weights (also called parameters). LLMs have a lot of weights. The Gemma 2 2B that is being used in this notebook has 2,617,270,528 parameters!\n",
        "\n",
        "Changing all of them can take quite some time and requires a lot of resources.\n",
        "\n",
        "To mitigate this issue, you are going to use a technique called: [LoRA: Low-Rank Adaptation](https://arxiv.org/abs/2106.09685)\n",
        "\n",
        "This technique, in summary, helps lower the number of trained weights needed by a lot, making fine-tuning more accessible.\n",
        "\n",
        "The key parameter used is the `rank`. In this notebook it set to 4 but you can use higher numbers to get better results but, of course, needed more resources.\n",
        "\n",
        "**TIP**: Train your model with lower ranks and evaluate the performance improvemnet on your task. Gradually increase the rank in subsequent trials and see if that further boosts performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RCucu6oHz53G"
      },
      "outputs": [],
      "source": [
        "# Enable LoRA for the model and set the LoRA rank to 4.\n",
        "gemma_lm.backbone.enable_lora(rank=lora_rank)\n",
        "gemma_lm.summary()\n",
        "\n",
        "# Limit the input sequence length (to control memory usage).\n",
        "gemma_lm.preprocessor.sequence_length = token_limit\n",
        "# Use AdamW (a common optimizer for transformer models).\n",
        "optimizer = keras.optimizers.AdamW(\n",
        "    learning_rate=lr_value,\n",
        "    weight_decay=0.01,\n",
        ")\n",
        "# Exclude layernorm and bias terms from decay.\n",
        "optimizer.exclude_from_weight_decay(var_names=[\"bias\", \"scale\"])\n",
        "\n",
        "gemma_lm.compile(\n",
        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    optimizer=optimizer,\n",
        "    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQQ47kcdpbZ9"
      },
      "source": [
        "Note that enabling LoRA reduces the number of trainable parameters significantly.\n",
        "\n",
        "From 2,617,270,528 to **2,928,640**\n",
        "\n",
        "To monitor the learning progress, you will evaluate the model at the end of each epoch and save the lora weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9G9sULJE5vpw"
      },
      "outputs": [],
      "source": [
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pc-Nea7L1fBc"
      },
      "outputs": [],
      "source": [
        "# !cd drive\n",
        "!mkdir -p ./drive/MyDrive/gemma_workshop\n",
        "# !cd .."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "26d9npFhAOSp"
      },
      "outputs": [],
      "source": [
        "class CustomCallback(keras.callbacks.Callback):\n",
        "  def on_epoch_end(self, epoch, logs=None):\n",
        "    model_name = f\"./drive/MyDrive/gemma_workshop/{lora_name}_{lora_rank}_epoch{epoch+1}.lora.h5\"\n",
        "    gemma_lm.backbone.save_lora_weights(model_name)\n",
        "\n",
        "    # Evaluate\n",
        "    text_gen(\"Write a title\")\n",
        "    text_gen(\"Write a poem\")\n",
        "\n",
        "history = gemma_lm.fit(train, epochs=train_epoch, batch_size=2, callbacks=[CustomCallback()])\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "plt.plot(history.history['loss'])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tn-jgVULyBXq"
      },
      "source": [
        "Note that the model began to grasp our intent more effectively from Epoch #3 onwards.\n",
        "\n",
        "To compare and contrast, it was utilized the \"Write a poem\" prompt. Interestingly, in Epoch #5, the model began to generate Portuguese in response to that prompt. This shift indicates a strong influence of our training dataset on the model's behavior. However, depending on your application, such a significant change might not be desirable. In such cases, Epoch #4 would be a more suitable choice."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P-tVAKmda2Zt"
      },
      "source": [
        "## Load LoRA\n",
        "\n",
        "Use the code below if you shared LoRA weights. It's much more lightweight than the model files themselves - for instance, a LoRA rank 4 weights file for a 10gb model might only be on the order of a few megabytes, easily shared over email."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hSW8-HRMa4ZB"
      },
      "outputs": [],
      "source": [
        "# Example Code for Load LoRA\n",
        "'''\n",
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(model_id)\n",
        "# Use the same LoRA rank that you trained\n",
        "gemma_lm.backbone.enable_lora(rank=4)\n",
        "train_epoch = 3\n",
        "# Load pre-trained LoRA weights\n",
        "gemma_lm.backbone.load_lora_weights(f\"/content/drive/MyDrive/{lora_name}_{lora_rank}_epoch{train_epoch}.lora.h5\")\n",
        "'''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ipg1u_wEKTxG"
      },
      "source": [
        "## Try a different sampler\n",
        "\n",
        "The top-K algorithm randomly picks the next token from the tokens of top K probability."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nV5mD_HqKZRF"
      },
      "outputs": [],
      "source": [
        "gemma_lm.compile(sampler=\"top_k\")\n",
        "text_gen(\"Write a title\")\n",
        "text_gen(\"Write a title\")\n",
        "text_gen(\"Write a title\")\n",
        "text_gen(\"Write a title\")\n",
        "text_gen(\"Write a title\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3m1XaCrlMu3Y"
      },
      "source": [
        "Try a slight different prompts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qC-MLxYWM1HU"
      },
      "outputs": [],
      "source": [
        "text_gen(\"Write a music title\")\n",
        "text_gen(\"Write a poem title\")\n",
        "text_gen(\"Write a blog title\")\n",
        "text_gen(\"Write a movie title\")\n",
        "text_gen(\"Write a novel title\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aEptDCED9tVp"
      },
      "source": [
        "## Publish your model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T3Qhrlyy5ReL"
      },
      "source": [
        "Lets save our model. It takes some time (~11 minutes) as it is a very large file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4TcvzBH995FE"
      },
      "outputs": [],
      "source": [
        "# Save the finetuned model\n",
        "my_model_name = \"my_gemma2_pt\"\n",
        "gemma_lm.save_to_preset(f\"./{my_model_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xQ4de1a79zy0"
      },
      "source": [
        "## Publishing on Hugging Face\n",
        "\n",
        "To publish your model on Hugging Face, you'll need to add your hugging face user (`HF_USER`) and an access token with write permission (`HF_TOKEN`) to the your secret keys."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rNj6rpB9LW4"
      },
      "outputs": [],
      "source": [
        "my_hf_username = userdata.get(\"HF_USER\")\n",
        "os.environ[\"HF_USER\"] = my_hf_username\n",
        "os.environ[\"HF_TOKEN\"] = userdata.get(\"HF_TOKEN\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H4v02q0H-QFL"
      },
      "outputs": [],
      "source": [
        "# Upload the model to Hugging Face Hub\n",
        "my_model_name = \"my_gemma2_pt\"\n",
        "hf_uri = f\"hf://{my_hf_username}/{my_model_name}\"\n",
        "keras_nlp.upload_preset(hf_uri, f\"./{my_model_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JzvyPgpX-XZk"
      },
      "source": [
        "## Publishing on Kaggle\n",
        "\n",
        "To publish on Kaggle you will need you username (already using if you followed the suggested setup process) and a name for your model and run the following cell\n",
        "\n",
        "**Note**: After the following cell, the model will be private to you on Kaggle. To share more broadly, change it to public."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H-T2XQtS-k3y"
      },
      "outputs": [],
      "source": [
        "# Upload the model to Kaggle\n",
        "my_user_name = os.environ[\"KAGGLE_USERNAME\"]\n",
        "my_model_name = \"my_gemma2_pt\"\n",
        "kaggle_uri = f\"kaggle://{my_user_name}/gemma/keras/{my_model_name}\"\n",
        "keras_nlp.upload_preset(kaggle_uri, f\"./{my_model_name}\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Workshop_How_to_Fine_tuning_Gemma.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
